//
// Created by shirley on 2024/12/23.
//

#include "DenseSolver.h"

#include <cmath>

DenseSolver::DenseSolver(const DenseMatrix& matrix, int thread, string function, string type): matrix(matrix), invMatrix(0,0)
{
    this->function = function;
    this->type = type;
    this->thread = thread;
}


vector<vector<double>>& DenseSolver::data_(shared_ptr<DenseMatrix>& matrix)
{
    return matrix->data;
}

vector<DenseMatrix>& DenseSolver::blocks_(shared_ptr<DenseMatrix>& matrix)
{
    return matrix->blocks;
}

// LU分解
bool DenseSolver::lU_Decomposition(shared_ptr<DenseMatrix>& L, shared_ptr<DenseMatrix>& U, const DenseMatrix& matrix)
{
    //获取矩阵natrix的行和列
    const auto rows = matrix.getRows();
    const auto cols = matrix.getCols();
    //matrix为非方阵
    if (rows != cols)
    {
        throw invalid_argument("LU decomposition requires a square matrix");
    }
    //初始化L、U矩阵
    L = make_shared<DenseMatrix>(rows, cols, 0.0);
    U = make_shared<DenseMatrix>(rows, cols);

    //获取矩阵A数据及L,U矩阵数据引用
    auto Matrix_Data = matrix.GetData();
    auto& L_Data = data_(L);
    auto& U_Data = data_(U);
    auto size = rows;

    //LU分解，填充L、U矩阵
    for (size_t i = 0; i < size; ++i)
    {
        // 检查对角线元素是否为零
        if (abs(Matrix_Data[i][i]) < 1e-10)
        {
            return false;
        }

        // 设置U矩阵的对角线元素
        U_Data[i][i] = Matrix_Data[i][i];

        // 设置U矩阵的第i行
#pragma omp parallel for schedule(static) num_threads(thread)
        for (int j = i + 1; j < size; ++j)
        {
            U_Data[i][j] = Matrix_Data[i][j];
        }

        // 设置L矩阵的对角线元素
        L_Data[i][i] = 1.0;

        // 计算L矩阵的第i列
#pragma omp parallel for schedule(static) num_threads(thread)
        for (int j = i + 1; j < size; ++j)
        {
            L_Data[j][i] = Matrix_Data[j][i] / U_Data[i][i];

            // 更新原始矩阵的第j行
            for (size_t k = i; k < size; ++k)
            {
                Matrix_Data[j][k] -= L_Data[j][i] * U_Data[i][k];
            }
        }
    }
    return true;
}

const DenseMatrix& DenseSolver::invMatrix_() const
{
    if (invMatrix.getRows() == 0 || invMatrix.getCols() == 0)
    {
        cout << "The matrix does not seem to be inverse at this time." << endl;
    }
    return invMatrix;
}

// 使用LU分解求逆
bool DenseSolver::LU_Inverse()
{
    //初始化L、U矩阵
    shared_ptr<DenseMatrix> L, U;
    //对矩阵matrix进行LU分解
    if (!lU_Decomposition(L, U, matrix))
    {
        printf("lU_Decomposition failed.\n");
        return false;
    }

    //获取矩阵matrix的行数
    const auto rows = matrix.getRows();
    //初始化matrix的逆矩阵invMatrix
    invMatrix = DenseMatrix(rows, rows);

    //LU分解求逆
    for (size_t i = 0; i < rows; ++i)
    {
        //初始化单元矩阵b
        DenseMatrix b(rows, 1, 0.0);
        b.setValue(i, 0, 1.0);

        //向前替换求解 Ly = b
        DenseMatrix y = forward_substitution(*L, b);

        //向后替换求解 Ux = y
        DenseMatrix x = backward_substitution(*U, y);

        //填充逆矩阵
        for (size_t j = 0; j < rows; ++j)
        {
            invMatrix.setValue(j, i, x.getValue(j, 0));
        }
    }
    return true;
}

bool DenseSolver::Inverse()
{
    if (type == "LU")
    {
        return LU_Inverse();
    }
    else
    {
        printf("Currently, only LU decomposition is supported！\n");
        return false;
    }
}

DenseMatrix DenseSolver::LU_Solve(DenseMatrix& B)
{
    //初始化L、U向量
    shared_ptr<DenseMatrix> L, U;
    //LU分解
    if (!lU_Decomposition(L, U, matrix))
    {
        cerr << "Matrix is singular and cannot be inverted." << endl;
        return {};
    }

    //获取矩阵matrix的行数
    const auto rows = matrix.getRows();
    //初始化matrix的逆矩阵invMatrix
    invMatrix = DenseMatrix(rows, rows);

    //利用LU分解求解线性方程组以求解插值系数
    auto coefficients = solveLinearSystem(L, U, B);

    return coefficients;
}

vector<DenseMatrix> DenseSolver::LU_Solve(vector<DenseMatrix>& B)
{
    //初始化L、U向量
    shared_ptr<DenseMatrix> L, U;
    //LU分解
    if (!lU_Decomposition(L, U, matrix))
    {
        cerr << "Matrix is singular and cannot be inverted." << endl;
        return {};
    }
    //获取矩阵matrix的行数
    const auto rows = matrix.getRows();
    //初始化matrix的逆矩阵invMatrix
    invMatrix = DenseMatrix(rows, rows);

    vector<DenseMatrix> coefficients(B.size());
    //利用LU分解求解线性方程组以求解插值系数
    for (size_t i = 0; i < B.size(); ++i)
    {
        coefficients[i] = solveLinearSystem(L, U, B[i]);
    }

    return coefficients;
}

DenseMatrix DenseSolver::solve(DenseMatrix& B)
{
    if (type == "LU")
    {
        return LU_Solve(B);
    }
    else
    {
        printf("Currently, only LU decomposition is supported！\n");
        return {};
    }
}

vector<DenseMatrix> DenseSolver::solve(vector<DenseMatrix>& B)
{
    if (type == "LU")
    {
        return LU_Solve(B);
    }
    else
    {
        printf("Currently, only LU decomposition is supported！\n");
        return {};
    }
}

DenseMatrix DenseSolver::solveLinearSystem(shared_ptr<DenseMatrix>& L, shared_ptr<DenseMatrix>& U, const DenseMatrix& b)
{
    //向前替换求解 Ly = b
    DenseMatrix y = forward_substitution(*L, b);
    //向后替换求解 Ux = y
    DenseMatrix x = backward_substitution(*U, y);

    return x;
}

DenseMatrix DenseSolver::forward_substitution(const DenseMatrix& L, const DenseMatrix& b)
{
    size_t n = L.getRows();
    DenseMatrix y(n, 1);

// #pragma omp parallel for schedule(static) num_threads(thread)
    for (size_t i = 0; i < n; ++i)
    {
        double sum = 0.0;
        for (size_t j = 0; j < i; ++j)
        {
            sum += L.getValue(i, j) * y.getValue(j, 0);
        }
        y.setValue(i, 0, (b.getValue(i, 0) - sum) / L.getValue(i, i));
    }
    return y;
}

DenseMatrix DenseSolver::backward_substitution(const DenseMatrix& U, const DenseMatrix& y)
{
    size_t n = U.getRows();
    DenseMatrix x(n, 1);
// #pragma omp parallel for schedule(static) num_threads(thread)
    for (int i = n - 1; i >= 0; --i)
    {
        double sum = 0.0;
        for (size_t j = i + 1; j < n; ++j)
        {
            sum += U.getValue(i, j) * x.getValue(j, 0);
        }
        x.setValue(i, 0, (y.getValue(i, 0) - sum) / U.getValue(i, i));
    }
    return x;
}


bool DenseSolver::blockLU_Decomposition(shared_ptr<DenseMatrix>& L, shared_ptr<DenseMatrix>& U, const DenseMatrix& matrix, size_t blockSize)
{
    const auto rows = matrix.getRows();
    const auto cols = matrix.getCols();

    if (rows != cols)
    {
        cout << "LU decomposition requires a square matrix" << endl;
        return false;
    }

    // 初始化L、U矩阵
    L = make_shared<DenseMatrix>(rows, cols, 0.0);
    U = make_shared<DenseMatrix>(rows, cols);

    auto& L_Data = data_((L));
    auto& U_Data = data_((U));
    auto Matrix_Data = matrix.GetData();
    for (size_t i = 0; i < rows; i += blockSize)
    {
        size_t brow = std::min(i + blockSize, rows); // 当前块的行数

        // 对角块
        shared_ptr<DenseMatrix> L_block;
        shared_ptr<DenseMatrix> U_block;
        DenseMatrix block(matrix, i, i, brow - i, brow - i); // 提取对角块

        lU_Decomposition(L_block, U_block, block);

        // 将分解得到的L_block和U_block复制到全局L和U矩阵中
        auto L_block_Data = L_block->GetData();
        auto U_block_Data = U_block->GetData();
        for (size_t k = i; k < brow; ++k)
        {
            for (size_t l = i; l < brow; ++l)
            {
                L_Data[k][l] = L_block_Data[k - i][l - i];
                U_Data[k][l] = U_block_Data[k - i][l - i];
            }
        }

        // 非对角块
        for (size_t j = brow; j < rows; j += blockSize)
        {
            size_t bcol = std::min(j + blockSize, rows); // 当前块的列数

            // 更新U矩阵的非对角块
            for (size_t k = i; k < brow; ++k)
            {
                for (size_t l = j; l < bcol; ++l)
                {
                    U_Data[k][l] = Matrix_Data[k][l];
                }
            }

            // 更新L矩阵的非对角块
            for (size_t k = j; k < bcol; ++k)
            {
                for (size_t l = i; l < brow; ++l)
                {
                    L_Data[k][l] = Matrix_Data[k][l] / U_Data[l][l];
                    for (size_t m = l; m < cols; ++m)
                    {
                        Matrix_Data[k][m] -= L_Data[k][l] * U_Data[l][m];
                    }
                }
            }
        }
    }

    return true;
}